"""Cookie池管理器

管理和维护各平台的Cookie池，支持自动轮换和失效检测
"""
import asyncio
import json
import logging
import random
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import redis.asyncio as aioredis
import aiofiles

logger = logging.getLogger(__name__)


class CookiePool:
    """Cookie池管理器"""
    
    def __init__(self, platform: str, redis_url: str = None, storage_path: str = "./cookies"):
        """初始化Cookie池
        
        Args:
            platform: 平台名称(weibo, zhihu, toutiao等)
            redis_url: Redis连接URL
            storage_path: 本地存储路径
        """
        self.platform = platform
        self.redis_url = redis_url
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(exist_ok=True)
        
        self.cookies: List[Dict] = []
        self.current_index = 0
        self.redis_client = None
        
        # Cookie健康状态
        self.cookie_stats: Dict[str, Dict] = {}
        
        # 配置
        self.max_failures = 3  # 最大失败次数
        self.check_interval = 3600  # 检查间隔（秒）
        
    async def initialize(self):
        """初始化连接"""
        if self.redis_url:
            try:
                self.redis_client = aioredis.from_url(
                    self.redis_url,
                    encoding="utf-8",
                    decode_responses=True
                )
                logger.info(f"{self.platform} Cookie池已连接到Redis")
            except Exception as e:
                logger.warning(f"Redis连接失败，使用本地存储: {e}")
        
        # 加载cookies
        await self.load_cookies()
    
    async def load_cookies(self):
        """加载Cookie池"""
        # 优先从Redis加载
        if self.redis_client:
            try:
                key = f"cookie_pool:{self.platform}"
                data = await self.redis_client.get(key)
                if data:
                    self.cookies = json.loads(data)
                    logger.info(f"从Redis加载了 {len(self.cookies)} 个{self.platform} cookies")
                    return
            except Exception as e:
                logger.error(f"从Redis加载失败: {e}")
        
        # 从本地文件加载
        cookie_file = self.storage_path / f"{self.platform}_cookies.json"
        if cookie_file.exists():
            try:
                async with aiofiles.open(cookie_file, 'r') as f:
                    content = await f.read()
                    self.cookies = json.loads(content)
                    logger.info(f"从文件加载了 {len(self.cookies)} 个{self.platform} cookies")
            except Exception as e:
                logger.error(f"从文件加载失败: {e}")
    
    async def save_cookies(self):
        """保存Cookie池"""
        # 保存到Redis
        if self.redis_client:
            try:
                key = f"cookie_pool:{self.platform}"
                await self.redis_client.set(
                    key,
                    json.dumps(self.cookies),
                    ex=86400  # 24小时过期
                )
            except Exception as e:
                logger.error(f"保存到Redis失败: {e}")
        
        # 保存到本地文件
        cookie_file = self.storage_path / f"{self.platform}_cookies.json"
        try:
            async with aiofiles.open(cookie_file, 'w') as f:
                await f.write(json.dumps(self.cookies, indent=2, ensure_ascii=False))
        except Exception as e:
            logger.error(f"保存到文件失败: {e}")
    
    def add_cookie(self, cookie: str, metadata: Dict = None):
        """添加Cookie
        
        Args:
            cookie: Cookie字符串
            metadata: Cookie元数据（用户信息等）
        """
        cookie_data = {
            'cookie': cookie,
            'add_time': datetime.now().isoformat(),
            'last_use': None,
            'use_count': 0,
            'fail_count': 0,
            'status': 'active',
            'metadata': metadata or {}
        }
        
        # 检查是否已存在
        for c in self.cookies:
            if c['cookie'] == cookie:
                logger.info(f"Cookie已存在，更新状态")
                c.update(cookie_data)
                return
        
        self.cookies.append(cookie_data)
        logger.info(f"添加新Cookie到{self.platform}池")
    
    def get_cookie(self) -> Optional[str]:
        """获取一个可用的Cookie
        
        Returns:
            Cookie字符串，如果没有可用的返回None
        """
        if not self.cookies:
            return None
        
        # 获取活跃的cookies
        active_cookies = [c for c in self.cookies if c['status'] == 'active']
        if not active_cookies:
            logger.warning(f"{self.platform}没有活跃的Cookie")
            return None
        
        # 轮询选择
        cookie_data = active_cookies[self.current_index % len(active_cookies)]
        self.current_index = (self.current_index + 1) % len(active_cookies)
        
        # 更新使用信息
        cookie_data['last_use'] = datetime.now().isoformat()
        cookie_data['use_count'] += 1
        
        return cookie_data['cookie']
    
    def get_random_cookie(self) -> Optional[str]:
        """随机获取一个可用的Cookie
        
        Returns:
            Cookie字符串
        """
        active_cookies = [c for c in self.cookies if c['status'] == 'active']
        if not active_cookies:
            return None
        
        cookie_data = random.choice(active_cookies)
        cookie_data['last_use'] = datetime.now().isoformat()
        cookie_data['use_count'] += 1
        
        return cookie_data['cookie']
    
    def mark_cookie_success(self, cookie: str):
        """标记Cookie使用成功
        
        Args:
            cookie: Cookie字符串
        """
        for c in self.cookies:
            if c['cookie'] == cookie:
                c['fail_count'] = 0
                c['last_success'] = datetime.now().isoformat()
                break
    
    def mark_cookie_failed(self, cookie: str, reason: str = None):
        """标记Cookie使用失败
        
        Args:
            cookie: Cookie字符串
            reason: 失败原因
        """
        for c in self.cookies:
            if c['cookie'] == cookie:
                c['fail_count'] += 1
                c['last_fail'] = datetime.now().isoformat()
                if reason:
                    c['fail_reason'] = reason
                
                # 超过最大失败次数，标记为失效
                if c['fail_count'] >= self.max_failures:
                    c['status'] = 'invalid'
                    logger.warning(f"{self.platform} Cookie失效: 连续失败{c['fail_count']}次")
                break
    
    async def validate_cookie(self, cookie: str, validator_func) -> bool:
        """验证Cookie有效性
        
        Args:
            cookie: Cookie字符串
            validator_func: 验证函数
            
        Returns:
            是否有效
        """
        try:
            is_valid = await validator_func(cookie)
            if is_valid:
                self.mark_cookie_success(cookie)
            else:
                self.mark_cookie_failed(cookie, "验证失败")
            return is_valid
        except Exception as e:
            self.mark_cookie_failed(cookie, str(e))
            return False
    
    async def refresh_pool(self):
        """刷新Cookie池，移除失效的Cookie"""
        before_count = len(self.cookies)
        
        # 移除失效的cookies
        self.cookies = [c for c in self.cookies if c['status'] != 'invalid']
        
        removed = before_count - len(self.cookies)
        if removed > 0:
            logger.info(f"移除了 {removed} 个失效的{self.platform} Cookie")
            await self.save_cookies()
    
    def get_stats(self) -> Dict:
        """获取Cookie池统计信息
        
        Returns:
            统计信息
        """
        total = len(self.cookies)
        active = len([c for c in self.cookies if c['status'] == 'active'])
        invalid = len([c for c in self.cookies if c['status'] == 'invalid'])
        
        return {
            'platform': self.platform,
            'total': total,
            'active': active,
            'invalid': invalid,
            'usage': {
                'total_uses': sum(c['use_count'] for c in self.cookies),
                'total_failures': sum(c['fail_count'] for c in self.cookies)
            }
        }
    
    async def close(self):
        """关闭连接"""
        await self.save_cookies()
        if self.redis_client:
            await self.redis_client.close()


class CookiePoolManager:
    """Cookie池管理器（管理多个平台）"""
    
    def __init__(self, redis_url: str = None):
        """初始化管理器
        
        Args:
            redis_url: Redis连接URL
        """
        self.redis_url = redis_url
        self.pools: Dict[str, CookiePool] = {}
    
    async def get_pool(self, platform: str) -> CookiePool:
        """获取指定平台的Cookie池
        
        Args:
            platform: 平台名称
            
        Returns:
            Cookie池实例
        """
        if platform not in self.pools:
            pool = CookiePool(platform, self.redis_url)
            await pool.initialize()
            self.pools[platform] = pool
        
        return self.pools[platform]
    
    async def add_cookie(self, platform: str, cookie: str, metadata: Dict = None):
        """添加Cookie到指定平台
        
        Args:
            platform: 平台名称
            cookie: Cookie字符串
            metadata: Cookie元数据
        """
        pool = await self.get_pool(platform)
        pool.add_cookie(cookie, metadata)
        await pool.save_cookies()
    
    async def get_cookie(self, platform: str) -> Optional[str]:
        """获取指定平台的Cookie
        
        Args:
            platform: 平台名称
            
        Returns:
            Cookie字符串
        """
        pool = await self.get_pool(platform)
        return pool.get_cookie()
    
    async def close_all(self):
        """关闭所有连接"""
        for pool in self.pools.values():
            await pool.close()


# 平台特定的Cookie验证器
async def validate_weibo_cookie(cookie: str) -> bool:
    """验证微博Cookie"""
    import aiohttp
    url = "https://weibo.com/ajax/profile/info"
    headers = {
        'Cookie': cookie,
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
    }
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers) as response:
                if response.status == 200:
                    data = await response.json()
                    return data.get('ok') == 1
    except:
        pass
    return False


async def validate_zhihu_cookie(cookie: str) -> bool:
    """验证知乎Cookie"""
    import aiohttp
    url = "https://www.zhihu.com/api/v4/me"
    headers = {
        'Cookie': cookie,
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
    }
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers) as response:
                return response.status == 200
    except:
        pass
    return False


async def main():
    """测试Cookie池"""
    manager = CookiePoolManager()
    
    # 添加测试Cookie（实际使用需要真实的Cookie）
    await manager.add_cookie('weibo', 'SUB=_2A...', {'username': 'test_user'})
    
    # 获取Cookie
    cookie = await manager.get_cookie('weibo')
    print(f"获取到Cookie: {cookie[:20]}..." if cookie else "没有可用Cookie")
    
    # 获取统计信息
    pool = await manager.get_pool('weibo')
    stats = pool.get_stats()
    print(f"Cookie池统计: {stats}")
    
    await manager.close_all()


if __name__ == "__main__":
    asyncio.run(main())